package com.fy.fyai.common.config;

import com.fy.fyai.model.entity.MessageContent;
import com.fy.fyai.service.AiChatConversationService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.context.annotation.Bean;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 自定义ChatMemory配置
 * 实现会话记忆功能，将聊天上下文保存到PostgreSQL数据库
 *
 * @author fwq
 * @date 2025-10-30
 */
@Slf4j
@Component
public class ChatMemoryService implements ChatMemory {

    @Resource
    private AiChatConversationService aiChatConversationService;

    /**
     * Bean方式提供ChatMemory实例
     */
    @Bean
    public ChatMemory chatMemory() {
        return this;
    }

    /**
     * 添加消息到PostgreSQL数据库
     *
     * @param conversationId 会话ID
     * @param messages       消息列表
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        aiChatConversationService.addMessageToConversation(conversationId , messages);
        log.info("成功保存{}条消息到PostgreSQL", messages.size());
    }

    /**
     * 从PostgreSQL数据库获取会话的历史消息
     *
     * @param conversationId 会话ID
     * @param lastN          获取最后N条消息
     * @return 消息列表，如果没有消息返回空列表（不能返回null）
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        log.info("接收到的请求：会话id: {}，获取最后 {} 条消息", conversationId, lastN);

        List<MessageContent> messageContentList = aiChatConversationService.getTopKChatMemory(conversationId, lastN);
        // 转换为Message列表（Message是抽象类，需要使用UserMessage或AssistantMessage）
        List<Message> messages = new ArrayList<>();
        for (MessageContent entity : messageContentList) {
            // 直接比较字符串，role是String类型
            if ("USER".equals(entity.getRole()) || "user".equalsIgnoreCase(entity.getRole())) {
                messages.add(new UserMessage(entity.getContent()));
            } else {
                messages.add(new AssistantMessage(entity.getContent()));
            }
        }

        log.info("成功从PostgreSQL获取{}条消息", messages.size());
        return messages;
    }

    /**
     * 清除会话的所有消息（从PostgreSQL删除）
     *
     * @param conversationId 会话ID
     */
    @Override
    public void clear(String conversationId) {
        aiChatConversationService.clearChatMemory(conversationId);
    }
}
